/-
Copyright (c) 2025 Marc Huisinga. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Marc Huisinga, Thomas R. Murrills
-/
module

public import Mathlib.Lean.Environment
public import Lean.Server.InfoUtils
public import Lean.Meta.TryThis
public import Batteries.Tactic.Lint.Misc
-- Import this linter explicitly to ensure that
-- this file has a valid copyright header and module docstring.
import Mathlib.Tactic.Linter.Header

/-!
# Additions to `Lean.Elab.InfoTree.Main`
-/

@[expose] public section

namespace Lean.Elab

open Lean.Meta
open Lean.Meta.Tactic.TryThis

/--
Collects all suggestions from all `TryThisInfo`s in `trees`.
`trees` is visited in order and suggestions in each tree are collected in post-order.
-/
def collectTryThisSuggestions (trees : PersistentArray InfoTree) :
    Array Suggestion :=
  go.run #[] |>.2
where
  /-- Visits all trees. -/
  go : StateM (Array Suggestion) Unit := do
    for tree in trees do
      tree.visitM' (postNode := visitNode)
  /-- Visits a node in a tree. -/
  @[nolint unusedArguments]
  visitNode (_ctx : ContextInfo) (i : Info) (_children : PersistentArray InfoTree) :
      StateM (Array Suggestion) Unit := do
    let .ofCustomInfo ci := i
      | return
    let some tti := ci.value.get? TryThisInfo
      | return
    modify (·.push tti.suggestion)

namespace InfoTree

/--
Finds the first result of `← f ctx info children` which is `some a`, descending the
tree from the top. Merges and updates contexts as it descends the tree.

`f` is **only** evaluated on nodes when some context is present. An initial context should be
provided via the `ctx?` argument if invoking `findSomeM?` during a larger traversal of the
infotree. A failure to provide `ctx? := some ctx` when `t` is not the outermost `InfoTree` is thus
likely to cause `findSomeM?` to always return `none`.
-/
partial def findSomeM? {m : Type → Type} [Monad m] {α}
    (f : ContextInfo → Info → PersistentArray InfoTree → m (Option α))
    (t : InfoTree) (ctx? : Option ContextInfo := none) : m (Option α) :=
  go ctx? t
where
  /-- Accumulates contexts and visits nodes if `ctx?` is not `none`. -/
  go ctx?
  | context ctx t => go (ctx.mergeIntoOuter? ctx?) t
  | node i ts => do
    let a ← match ctx? with
      | none => pure none
      | some ctx => f ctx i ts
    match a with
    | some a => pure a
    | none => ts.findSomeM? (go <| i.updateContext? ctx?)
  | hole _ => pure none

/--
Finds the first result of `f ctx info children` which is `some a`, descending the
tree from the top. Merges and updates contexts as it descends the tree.

`f` is **only** evaluated on nodes when some context is present. An initial context should be
provided via the `ctx?` argument if invoking `findSome?` during a larger traversal of the infotree.
A failure to provide `ctx? := some ctx` when `t` is not the outermost `InfoTree` is thus likely to
cause `findSome?` to always return `none`.
-/
def findSome? {α} (f : ContextInfo → Info → PersistentArray InfoTree → Option α)
    (t : InfoTree) (ctx? : Option ContextInfo := none) : Option α :=
  Id.run <| t.findSomeM? f ctx?

/--
Returns the value of `f ctx info children` on the outermost `.node info children` which has
context, having merged and updated contexts appropriately.

If `ctx?` is `some ctx`, `ctx` is used as an initial context. A `ctx?` of `none` should **only** be
used when operating on the first node of the entire infotree. Otherwise, it is likely that no
context will be found.
-/
def onHighestNode? {α} (t : InfoTree) (ctx? : Option ContextInfo)
    (f : ContextInfo → Info → PersistentArray InfoTree → α) : Option α :=
  t.findSome? (ctx? := ctx?) fun ctx i ch => some (f ctx i ch)

/--
Get the `parentDecl`s of every elaborated body.

This includes `let rec`/`where` definitions, but excludes decls without "bodies" (such as
`alias`es, `structure`s, declarations generated by attributes like `@[ext]`, and so on) as we
might find by considering every `parentDeclCtx` throughout the infotree.

Assumes that every body elaboration proceeds through `Lean.Elab.Term.BodyInfo`.
-/
def getDeclsByBody (t : InfoTree) : List Name :=
  t.collectNodesBottomUp fun ctx i _ decls =>
    match i with
    | .ofCustomInfo i =>
      if i.value.typeName == ``Lean.Elab.Term.BodyInfo then
        if let some decl := ctx.parentDecl? then
          decl :: decls
        else decls
      else decls
    | _ => decls

/--
Get the declarations elaborated in the infotree `t` which are theorems according to the
environment. This includes e.g. `instance`s of `Prop` classes in addition to declarations declared
using the keyword `theorem` directly.
-/
def getTheorems (t : InfoTree) (env : Environment) : List ConstantVal :=
  t.getDeclsByBody.filterMap env.findTheoremConstVal?

end Lean.Elab.InfoTree
